#include <fstream>
#include <iostream>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#include <assert.h>
#include <stdlib.h>
#include <boost/progress.hpp>
#include <boost/tokenizer.hpp>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#include <LoadMatrix.h>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
bool str2bool( const std::string & s ) {
   if( s == "True"  ) { return true;  }
   if( s == "False" ) { return false; }
   assert(0);
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ExtractRawFileName( const std::string & fileName, std::string & rawFileName ) {

   std::string t1 = "";
   std::string t2 = "";
   boost::char_separator<char> sep("./");
   boost::tokenizer<boost::char_separator<char> > tokens(fileName, sep);
   for( boost::tokenizer<boost::char_separator<char> >::iterator it = tokens.begin(); it != tokens.end(); ++it ) {
      DEBUG_PRINT( "token = " << *it );
      // keep two most recent tokens in t1 & t2
      t2 = t1;
      t1 = *it;
      // if the most recent token is "mtx" then the raw file name is found
      if( t1 == "mtx" ) {
         break;
      }
   }
   rawFileName = t2;

}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void CheckSupportedMtxTypes( const MatrixProperties & mtxInfo ) {
  if( (mtxInfo.matrix_type != "array") && (mtxInfo.matrix_type != "coordinate") ) {
    std::cout << "Unknown matrix type: " << mtxInfo.matrix_type << std::endl;
    exit(1);
  }
  else {
    if( mtxInfo.matrix_type == "array" ) {
      std::cout << "Dense matrix type (array) is not supported yet\n";
      exit(1);
    }
  }
  if( (mtxInfo.element_type != "real") && (mtxInfo.element_type != "integer") && (mtxInfo.element_type != "complex") && (mtxInfo.element_type != "pattern") ) {
	 std::cout << "Unknown element type: " << mtxInfo.element_type << std::endl;
    exit(1);
  }
  //else {
  //  if( mtxInfo.element_type == "complex" ) {
  //    std::cout << "Matrix with complex elements is not supported yet\n";
  //    exit(1);
  //  }
  //}
  if( (mtxInfo.symmetry_type != "general") && (mtxInfo.symmetry_type != "symmetric") && (mtxInfo.symmetry_type != "skew-symmetric") && (mtxInfo.symmetry_type != "Hermitian") ) {
    std::cout << "Unknown symmetry type: " << mtxInfo.symmetry_type << std::endl;
    exit(1);
  }
  else {
    if( mtxInfo.symmetry_type == "Hermitian" ) {
      std::cout << "Hermitian matrix is not supported yet\n";
      exit(1);
    }
  }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ProcessMMFileHeader(const std::string & fileName, std::ifstream & inputFile, MatrixProperties & mtxInfo ) {
   mtxInfo.fullFileName    = fileName;
   mtxInfo.isSymmetric     = false;
   mtxInfo.isSkewSymmetric = false;
   mtxInfo.isComplex       = false;
   mtxInfo.isReal          = false;
   mtxInfo.isInteger       = false;
   mtxInfo.isPattern       = false;
      
   mtxInfo.nnz             = 0ULL;
   mtxInfo.nnzOffDiag      = 0ULL;
   mtxInfo.diagNonZeros    = 0ULL;
   mtxInfo.storedNonZeros  = 0ULL;
   
   ExtractRawFileName( mtxInfo.fullFileName, mtxInfo.rawFileName );
   
   std::string line;
   if( std::getline( inputFile, line )) {
      boost::tokenizer<> tok( line );
      int i=0;
      for( boost::tokenizer<>::iterator itr=tok.begin(); itr != tok.end(); ++itr ) {
         switch (i) {
            case 0: 
               if ( *itr != "MatrixMarket" ) {
                 std::cout << "File format is not Matrix Market" << std::endl;
                 exit( -1 );
               }
               break;
            case 1:
               if ( *itr != "matrix" ) {
                 std::cout << "File does not contain a matrix" << std::endl;
                 exit( -1 );
               }
               break;
            case 2:
               mtxInfo.matrix_type = *itr;
               break;
            case 3: 
               mtxInfo.element_type = *itr;
               break;
            case 4:
               mtxInfo.symmetry_type = *itr;
               break;
            default:
               std::cout << "Header format is not correct" << std::endl;
               exit( -1 );
               break;
         }
         i++;
      }
   }
   CheckSupportedMtxTypes(mtxInfo);
   // Skip over comments
   std::getline( inputFile, line );
   while( line[0] == '%' ) {
     DEBUG_PRINT( "mtx file header line = " << line );
     std::getline( inputFile, line );
   }
   
   // Read dimensions
   U64 rows;
   U64 columns;
   U64 elements;
   
   #ifdef __MAC_OSX__
   sscanf( line.c_str(), "%lld %lld %lld\n", &rows, &columns, &elements );
   #else
   sscanf( line.c_str(), "%ld %ld %ld\n",    &rows, &columns, &elements );
   #endif
   mtxInfo.m    = rows;
   mtxInfo.n    = columns;
   mtxInfo.nnz  = elements;

   if( mtxInfo.element_type ==  "real"      ) { mtxInfo.isReal      = true;  }
   if( mtxInfo.element_type ==  "integer"   ) { mtxInfo.isInteger   = true;  }
   if( mtxInfo.element_type ==  "pattern"   ) { mtxInfo.isPattern   = true;  }
   if( mtxInfo.element_type ==  "complex"   ) { mtxInfo.isComplex   = true;  }
   
   if( mtxInfo.symmetry_type == "symmetric" )      { mtxInfo.isSymmetric     = true;  }
   if( mtxInfo.symmetry_type == "skew_symmetric" ) { mtxInfo.isSkewSymmetric = true;  }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template<typename T> void PrintMappedMtxToScreen( std::map<Coordinate, T> & A, const std::string & mtxName ) {
  typename std::map< Coordinate, T>::iterator it;
  std::cout << "Printing matrix " << mtxName << std::endl;
  for( it = A.begin(); it != A.end(); it++ ) {
    Coordinate c = it->first;
    T          v = it->second;
    std::cout << mtxName << "[" << c.row << ", " << c.col << "] = " << v << std::endl;
  }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template<typename T> void MoveLowerHalfToUpperHalfKernel( std::map<Coordinate, T> & A ) {
   
   std::map< Coordinate, T > buf;
   typename std::map< Coordinate, T >::iterator it;
   for( it = A.begin(); it != A.end(); it++ ) {
      Coordinate c = it->first;
      T          v = it->second;
      c.Swap();
      buf[c] = v;
   }
   A.clear();
   A = buf;
   
}
void MoveLowerHalfToUpperHalf( std::map<Coordinate, S64>                   & A ) { MoveLowerHalfToUpperHalfKernel<S64>                  ( A ); }
void MoveLowerHalfToUpperHalf( std::map<Coordinate, double>                & A ) { MoveLowerHalfToUpperHalfKernel<double>               ( A ); }
void MoveLowerHalfToUpperHalf( std::map<Coordinate, std::complex<double> > & A ) { MoveLowerHalfToUpperHalfKernel<std::complex<double> >( A ); }
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void CopyBufferToComplexValuedMtx( void * vpA, void * vpB ) {

   std::map< Coordinate, std::complex<double> > * pA = (std::map< Coordinate, std::complex<double> > *)vpA;
   std::map< Coordinate, std::complex<double> > * pB = (std::map< Coordinate, std::complex<double> > *)vpB;
   
   std::map< Coordinate, std::complex<double> >::iterator it;
   
   for( it = pB->begin(); it != pB->end(); it++ ) {
      Coordinate           c = it->first;
      std::complex<double> v = it->second;
      (*pA)[ c ] = v;
   }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template<typename T> void ReadInMMElementsKernel(std::map<Coordinate, T> & A, std::ifstream & inputFile, MatrixProperties & mtxInfo, const bool expandSymmMats ) {

   const bool hasSymmetry = mtxInfo.isSymmetric || mtxInfo.isSkewSymmetric;
   
   Coordinate c;
   std::map<Coordinate, std::complex<double> > B;
   std::string scanfFmt;
   std::string line;
   
   U64 numNonZeroElements = 0;
   U64 diagNonZeros       = 0;
   U64 storedNonZeros     = 0;
   bool readingUpperHalf  = true;
   bool firstOffDiagElement = true;
   
   const std::complex<double> zero(0.0,0.0);
   const std::string scanfPattFmt = "%ld %ld\n";
   const std::string scanfRealFmt = "%ld %ld %lf\n";
   const std::string scanfIntsFmt = "%ld %ld %ld\n";
   const std::string scanfImagFmt = "%ld %ld %lf %lf\n";
   
   
   //----- Read in elements      
   while( std::getline(inputFile,line) ) {
      U64 i, j;
      if( mtxInfo.isReal ) {
         double a;
         sscanf( line.c_str(), scanfRealFmt.c_str(), &i, &j, &a );
         if( a == 0.0 ) { continue; }
         c.row = i-1;
         c.col = j-1;
         A[c] = a;
         numNonZeroElements++;
         storedNonZeros++;
         if( i == j ) {
            diagNonZeros++;
         }
         else {
            if( hasSymmetry ) {
               if( expandSymmMats ) {
                  c.row = j-1;
                  c.col = i-1;
                  if( mtxInfo.isSkewSymmetric ) { A[c] = (long int)(-1.0 * a);  }
                  else                          { A[c] = (long int)( 1.0 * a);  }
                  storedNonZeros++;
               }
            }
         }
      }
      else if( mtxInfo.isInteger ) {
         S64 a;
         sscanf( line.c_str(), scanfIntsFmt.c_str(), &i, &j, &a );
         if( a == 0.0 ) { continue; }
         c.row = i-1;
         c.col = j-1;
         A[c] = a;
         numNonZeroElements++;
         storedNonZeros++;
         if( i == j ) {
            diagNonZeros++;
         }
         else {
            if( hasSymmetry ) {
               if( expandSymmMats ) {
                  c.row = j-1;
                  c.col = i-1;
                  if( mtxInfo.isSkewSymmetric ) { A[c] = -1LL * a;  }
                  else                          { A[c] = a;         }
                  storedNonZeros++;
               }
            }
         }
      }
      else if( mtxInfo.isPattern ) {
         sscanf( line.c_str(), scanfPattFmt.c_str(), &i, &j );
         c.row = i-1;
         c.col = j-1;
         A[c] = 1;
         numNonZeroElements++;
         storedNonZeros++;
         if( i == j ) {
            diagNonZeros++;
         }
         else {
            if( hasSymmetry ) {
               if( expandSymmMats ) {
                  c.row = j-1;
                  c.col = i-1;
                  A[c] = 1;
                  storedNonZeros++;
               }
            }
         }
      }
      else if( mtxInfo.isComplex ) {
         double rv;
         double iv;
         sscanf( line.c_str(), scanfImagFmt.c_str(), &i, &j, &rv, &iv );
         std::complex<double> v(rv, iv);
         if( v == zero ) { continue; }
         c.row = i-1;
         c.col = j-1;
         B[c] = v;
         numNonZeroElements++;
         storedNonZeros++;
         if( i == j ) {
            diagNonZeros++;
         }
         else {
            if( hasSymmetry ) {
               if( expandSymmMats ) {
                  c.row = j-1;
                  c.col = i-1;
                  if( mtxInfo.isSkewSymmetric ) { B[c] = -1.0 * v;  }
                  else                          { B[c] = v;         }
                  storedNonZeros++;
               }
            }
         }      
      }
      if( hasSymmetry ) {
         if( i != j ) {
            if( firstOffDiagElement ) {
               firstOffDiagElement = false;
               if( i < j ) { readingUpperHalf = true;  }
               else        { readingUpperHalf = false; }
            }
            else {
               if( i < j ) { assert(  readingUpperHalf ); }
               else        { assert( !readingUpperHalf ); }
            }
         }
      }
   }
   if( mtxInfo.isComplex ) { CopyBufferToComplexValuedMtx( &A, &B ); }
   if( hasSymmetry && !expandSymmMats && !readingUpperHalf ) {
      DEBUG_PRINT("Moving lower half to upper half because:");
      DEBUG_PRINT("Matrix symmetry detected, and lower half stored in mtx market file, and expandSymmMats not set.");
      MoveLowerHalfToUpperHalf( A );
      //if( mtxInfo.isReal    ) { MoveLowerHalfToUpperHalf( A ); }
      //if( mtxInfo.isInteger ) { MoveLowerHalfToUpperHalf( A ); }
      //if( mtxInfo.isComplex ) { MoveLowerHalfToUpperHalf( A ); }
   }
   if(mtxInfo.nnz != numNonZeroElements) {
      std::cout << mtxInfo.nnz << " nnz recorded in file header, but " << numNonZeroElements << " actually read in." << std::endl;
      mtxInfo.nnz = numNonZeroElements;
   }
   mtxInfo.diagNonZeros    = diagNonZeros;
   mtxInfo.storedNonZeros  = storedNonZeros;
   mtxInfo.nnzOffDiag = numNonZeroElements - diagNonZeros;
}
void ReadInMMElements( std::map<Coordinate, S64>                   & A, std::ifstream & inputFile, MatrixProperties & mtxInfo, const bool expandSymmMats ) {
   ReadInMMElementsKernel< S64 >( A, inputFile, mtxInfo, expandSymmMats );
}
void ReadInMMElements( std::map<Coordinate, double>                & A, std::ifstream & inputFile, MatrixProperties & mtxInfo, const bool expandSymmMats ) {
   ReadInMMElementsKernel< double >( A, inputFile, mtxInfo, expandSymmMats );
}
void ReadInMMElements( std::map<Coordinate, std::complex<double> > & A, std::ifstream & inputFile, MatrixProperties & mtxInfo, const bool expandSymmMats ) {
   ReadInMMElementsKernel< std::complex<double> >( A, inputFile, mtxInfo, expandSymmMats );
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void MtxMarket2SparseArray( const std::string & mtxFileName, const std::string & inFmt, std::map< U64, double > & spa, MatrixProperties & mtxMetaData ) {

   DEBUG_PRINT( "MtxMarket2SparseArray()" );

   std::map< Coordinate, double > mtx;

   std::ifstream inputFile( mtxFileName.c_str() );

   ProcessMMFileHeader( mtxFileName, inputFile, mtxMetaData );

   std::string fmt = inFmt;
   if( inFmt == "CSR" ) { fmt = "CSR"; }
   
   U64 nnz = 0;
   U64 m = RoundUpToPow2( mtxMetaData.m );
   U64 n = RoundUpToPow2( mtxMetaData.n );
   U64 r;
   U64 c;
   U64 p;
   double v;
   assert( !mtxMetaData.isComplex );
   assert(  mtxMetaData.isReal    );

   std::string line;
   const std::string scanfRealFmt = "%ld %ld %lf\n";
   
   if( mtxMetaData.isSymmetric ) {
      while( std::getline( inputFile, line ) ) {
      
         sscanf( line.c_str(), scanfRealFmt.c_str(), &r, &c, &v );
         if( v == 0.0 ) { continue; }
         
         p = RC2Pos( r-1, c-1, m, n, fmt );
         spa[ p ] = v;
         
         p = RC2Pos( c-1, r-1, m, n, fmt );
         spa[ p ] = v;
         
         // DEBUG_PRINT( "... nnzIdx = " << nnz << ", r = " << r << ", c = " << c << ", v = " << v << ", pos = " << p );
         nnz++;
         if( r != c ) { nnz++; }
      }
   }
   else {
      while( std::getline( inputFile, line ) ) {
      
         sscanf( line.c_str(), scanfRealFmt.c_str(), &r, &c, &v );
         if( v == 0.0 ) { continue; }
         
         p = RC2Pos( r-1, c-1, m, n, fmt );
      
         spa[ p ] = v;
         
         // DEBUG_PRINT( "... nnzIdx = " << nnz << ", r = " << r << ", c = " << c << ", v = " << v << ", pos = " << p );
         nnz++;
      }
   }
   
   assert( nnz == spa.size() );
   if( nnz != mtxMetaData.nnz ) {
      std::cout << mtxMetaData.nnz << " nnz indicated in mtx file but " << nnz << " read in." << std::endl;
      mtxMetaData.nnz = nnz;
   }
}
